package LibDL.eval.cluster;

import LibDL.core.*;
import LibDL.eval.SupervisedEvaluator;

public abstract class AbstractSupervisedClusterEvaluator implements SupervisedEvaluator {
    private Tensor truth;
    private Tensor pred;

    protected abstract double core(Tensor truth, Tensor pred);

    @Override
    public double evaluate(Tensor truth, Tensor pred) {
        setPred(pred);
        setTruth(truth);
        return core(getTruth(),getPred());
    }

    @Override
    public double evaluate() {
        return core(getTruth(),getPred());
    }

    public Tensor getContingencyMatrix(){
        StdVectorTensor classVec = functional.unique(getTruth(),true,true,false);
        StdVectorTensor clusterVec = functional.unique(getPred(),true,true,false);
        Tensor classes = classVec.get(0);
        StdVectorLong classIdx = classVec.get(1).tolist_long();
        Tensor clusters = clusterVec.get(0);
        StdVectorLong clusterIdx = clusterVec.get(1).tolist_long();

        Tensor contingency = functional.zeros(new StdVectorLong(new int[]{classes.size(0),clusters.size(0)}))
                .to(Dtype.FLOAT64);
        if (classIdx.size() != clusterIdx.size()){
            throw new IllegalArgumentException("The length of pred and truth must be the same.");
        }
        for (int i = 0;i<classIdx.size();i++){
            Tensor row = functional.zeros(contingency.size(1)).to(Dtype.FLOAT64)
                    .index_put_(functional.as_tensor(clusterIdx.get(i)),functional.as_tensor(1).to(Dtype.FLOAT64))
                    .unsqueeze(0);
            contingency.index_add_(0,functional.as_tensor(classIdx.get(i)),row);
        }
        return contingency;
    }
    protected long factorial(int n){
        return n>1? n*factorial(n-1) : 1;
    }
    protected long combination(int n, int m) {
        return (n >= m) ? factorial(n) / factorial(n - m) / factorial(m) : 0;
    }
    protected long comb2(int n){
        return combination(n,2);
    }
    protected void checkClustering(Tensor truth,Tensor pred){
        if (truth.sizes().size() != 1 || pred.sizes().size() != 1){
            throw new IllegalArgumentException("Truth and pred must be 1D. Found:" +
                    " truth.sizes().size() == "+truth.sizes().size()+
                    " pred.sizes().size() == "+pred.sizes().size());
        }
        if (truth.size(0) != pred.size(0)) {
            throw new IllegalArgumentException("The lengths of truth and pred must be consistent.");
        }
    }

    public Tensor getTruth() {
        return truth;
    }

    public AbstractSupervisedClusterEvaluator setTruth(Tensor truth) {
        this.truth = truth;
        return this;
    }

    public Tensor getPred() {
        return pred;
    }

    public AbstractSupervisedClusterEvaluator setPred(Tensor pred) {
        this.pred = pred;
        return this;
    }
}
